package com.leetcode.algs4.graph.mst;

/**
 * @author Dennis Li
 * @date 2020/6/18 19:40
 */
public class PrimMST {
    // 距离树最近的边
    private Edge[] edgeTo;
    private double[] distTo;
    private boolean[] marked;
    // 保存有效的横切边
    private IndexMinPQ<Double> pq;

    public PrimMST(EdgeWeightGraph G) {
        edgeTo = new Edge[G.V()];
        distTo = new double[G.V()];
        marked = new boolean[G.V()];
        for (int v = 0; v < G.V(); v++) {
            distTo[v] = Double.POSITIVE_INFINITY;
        }
        pq = new IndexMinPQ<>(G.V());
        distTo[0] = 0.0;
        pq.insert(0, 0.0);
        while (!pq.isEmpty())
            // 将最近的顶点添加到树中
            visit(G, pq.delMin());
    }

    private void visit(EdgeWeightGraph G, int v) {
        marked[v] = true;
        for (Edge e : G.adj(v)) {
            int w = e.other(v);
            if (marked[w]) continue; // v-w 失效，已经加入到mst中
            // 权值更小
            if (e.weight() < distTo[w]) {
                // 更新最小边和对应的权值
                edgeTo[w] = e;
                distTo[w] = e.weight();
                // 如果pq中已经有w，那么将它对应的value改为distTo[w]
                if (pq.contains(w)) pq.changeKey(w, distTo[w]);
                // 否则将它加入到横切边集合中
                else pq.insert(w, distTo[w]);
            }
        }
    }
}
